import os

# os.environ['OPENAI_API_KEY'] = "sk"
# os.environ['OPENAI_API_BASE'] = "https://"
# os.environ['HUGGINGFACEHUB_API_TOKEN'] = ''
from dotenv import load_dotenv

load_dotenv()

import streamlit as st
from PyPDF2 import PdfReader
import fitz  # PyMuPDF

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI

# 设置Streamlit应用标题
st.title("PDF问答应用")

# 上传PDF文件
uploaded_file = st.file_uploader("选择一个PDF文件", type="pdf")

if uploaded_file is not None:
    try:
        pdf_text = ""
        pdf = fitz.open(stream=uploaded_file.read(), filetype="pdf")  # 使用流打开PDF
        for page in pdf:
            pdf_text += page.get_text()  # 提取文本
        pdf.close()

        if not pdf_text:
            st.error("PDF 文件中未找到任何有效文本。")
        else:
            st.success("成功提取文本。")
            st.write(pdf_text[:500] + "...")  # 显示前500个字符

            # 文本分割
            # text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
            texts = text_splitter.split_text(pdf_text)

            # 确保文本分割结果非空
            if not texts:
                st.error("分割后的文本为空，请检查 PDF 文件的内容。")
            else:
                st.success(f"成功分割文本，分割出的文本块数量: {len(texts)}")

                # 创建嵌入
                embeddings = OpenAIEmbeddings()

                # 使用 Chroma 创建向量存储
                vectorstore = Chroma.from_texts(texts, embeddings, persist_directory="./chroma_db")  # 设置持久化目录

                # 创建问答链
                qa_chain = RetrievalQA.from_chain_type(
                    llm=OpenAI(),
                    chain_type="stuff",
                    retriever=vectorstore.as_retriever()
                )

                # 用户输入问题
                user_question = st.text_input("请输入您的问题：")

                if user_question:
                    # 获取答案
                    answer = qa_chain.run(user_question)
                    st.write("回答：", answer)
    except Exception as e:
        st.error(f"发生错误: {e}")  # 显示错误信息
